import datetime
import os

import torch.cuda
from flask import Flask, Blueprint, request, jsonify
from .train import start_train
from app.admin.models import MlRoute, TModules
from mlwrite import db
from mlwrite.model import JsonResult
from app.admin.log import red_logs, logging_

json_result = JsonResult(0, 'success', None)

bp = Blueprint(name='train', import_name=__name__, url_prefix='/train')


@bp.route('/start', methods=['POST'])
def train():
    label_id = request.form.get('label_id')
    module = TModules.query.filter_by(module_code=label_id).first()
    module.status = 2
    db.session.commit()
    try:
        start_train(module)
    except RuntimeError:
        module.status = 1
        db.session.commit()
    finally:
        torch.cuda.empty_cache()
    return jsonify(json_result.success(None))
